from keras.models import Sequential
from keras.layers import Dense, Activation, Flatten, Reshape
from keras.layers import Conv2D, Conv2DTranspose, UpSampling2D, MaxPooling2D
from keras.layers import LeakyReLU, Dropout
from keras.layers import BatchNormalization
from keras.optimizers import Adam, RMSprop,Nadam


class DCGAN(object):

    def __init__(self, img_rows, img_cols, channel):
        self.img_rows = img_rows
        self.img_cols = img_cols
        self.channel = channel
        self.D = None
        self.G = None
        self.AM = None
        self.DM = None

    def discriminator(self):
        if self.D:
            return self.D
        self.D = Sequential()
        depth = 32
        dropout = 0.4
        input_shape = (self.img_rows, self.img_cols, self.channel)
        self.D.add(Conv2D(depth * 1, 5, strides=2, input_shape=input_shape,
                          padding='same', activation=LeakyReLU(alpha=0.2)))
        self.D.add(Dropout(dropout))

        self.D.add(Conv2D(depth * 2, 5, strides=2, padding='same',
                          activation=LeakyReLU(alpha=0.2)))
        self.D.add(MaxPooling2D(pool_size=(2, 2)))
        self.D.add(Dropout(dropout))

        self.D.add(Conv2D(depth * 4, 5, strides=2, padding='same',
                          activation=LeakyReLU(alpha=0.2)))
        self.D.add(MaxPooling2D(pool_size=(2, 2)))
        self.D.add(Dropout(dropout))

        self.D.add(Conv2D(depth * 8, 5, strides=1, padding='same',
                          activation=LeakyReLU(alpha=0.2)))
        self.D.add(Dropout(dropout))

        self.D.add(Flatten())
        self.D.add(Dense(64))
        self.D.add(Activation('relu'))
        self.D.add(Dropout(dropout))
        self.D.add(Dense(1))
        self.D.add(Activation('sigmoid'))
        return self.D

    def generator(self):
        if self.G:
            return self.G
        self.G = Sequential()
        dropout = 0.4
        depth = 256
        dim_x = round(self.img_rows / 4)
        dim_y = round(self.img_cols / 4)
        self.G.add(Dense(depth, input_dim=100))
        self.G.add(Activation('relu'))
        self.G.add(Dense(dim_x * dim_y * depth))
        self.G.add(BatchNormalization(momentum=0.9))
        self.G.add(Activation('relu'))
        self.G.add(Reshape((dim_x, dim_y, depth)))
        self.G.add(Dropout(dropout))

        self.G.add(UpSampling2D())
        self.G.add(Conv2DTranspose(int(depth / 2), 5, padding='same'))
        self.G.add(BatchNormalization(momentum=0.9))
        self.G.add(Activation('relu'))

        self.G.add(UpSampling2D())
        self.G.add(Conv2DTranspose(int(depth / 4), 5, padding='same'))
        self.G.add(BatchNormalization(momentum=0.9))
        self.G.add(Activation('relu'))

        self.G.add(Conv2DTranspose(int(depth / 8), 5, padding='same'))
        self.G.add(BatchNormalization(momentum=0.9))
        self.G.add(Activation('relu'))

        self.G.add(Conv2DTranspose(self.channel, 5, padding='same'))
        self.G.add(Activation('sigmoid'))
        return self.G

    def discriminator_model(self):
        if self.DM:
            return self.DM
        # optimizer = RMSprop(lr=0.0008, clipvalue=1.0, decay=6e-8)
        optimizer = Nadam(lr=0.0008)
        self.DM = Sequential()
        self.DM.add(self.discriminator())
        self.DM.compile(optimizer=optimizer,
                        loss='binary_crossentropy', metrics=['accuracy'])
        return self.DM

    def adversarial_model(self):
        if self.AM:
            return self.AM
        optimizer = Nadam(lr=0.0008)
        # optimizer = RMSprop(lr=0.0006, clipvalue=1.0, decay=3e-8)
        self.AM = Sequential()
        self.AM.add(self.generator())
        self.discriminator().trainable = False
        self.AM.add(self.discriminator())
        self.AM.compile(optimizer=optimizer,
                        loss='binary_crossentropy', metrics=['accuracy'])
        self.discriminator().trainable = True
        return self.AM
